import logging
import numpy as np

from util.util_log import setup_logging
from util.util_csv import save_csv, save_to_csv
from util.util_image import save_img, save_img_xyz, save_line_chart, save_images_3x2, draw_img
from util.util_ris_pattern import point_2_phase, phase_2_pattern, phase_2_pattern_xyz, eps, \
    phase_2_bit, phaseBit_2_deg, phase_2_pattern_xyz_fft

from multi_beam_trace.util_create_trace import create_scan_traces_theta, \
    create_scan_traces_phi, create_scan_traces_theta_phi
from multi_beam_trace.util_phase_pattern import traces_2_phaseBits, traces_2_phaseRads, phaseBits_2_patterns
from multi_beam_trace.util_analysis import get_phase_neighbor_diff, get_antenna_indicators_normalized, \
    get_antenna_indicators_normalized_item, compare_arrays
from multi_beam_trace.FitnessMultiBeamTrace import phaseBit_2_pattern, pattern_2_peaks
from multi_beam_trace.BPSO import BPSO
from multi_beam_trace.FitnessMultiBeamTrace import FitnessMultiBeamTrace

from multi_beam.multi_beam_PS_complex import psm_complex_beam_n
from multi_beam.multi_beam_QE import qe_beam_N


# 配置日志，默认打印到控制台，也可以设置打印到文件
# setup_logging()
setup_logging(log_file="../files/multi-beam-trace/beam2/multi-group-BPSO/trace_2025-04-04-0903-16beam.txt")
# 获取日志记录器并记录日志
logger = logging.getLogger("[multi-beam-trace-BPSO]")


# ============================================= 天线相关计算 =======================================
# 码阵转bit
def phases_rad_2_bit(phaseRads, bit_num):
    phaseBits = []
    for phaseRad in phaseRads:
        phaseBit, phaseDeg = phase_2_bit(phaseRad, bit_num)
        phaseBits.append(phaseBit)
    return phaseBits


# ============================================= 主方法 ===========================================
# 多波束合成: PS-complex
def multi_beam_ps_complex(phaseRads_list, bit_num):
    # 整理数据, 将所有[0]放第一个, [1]放第二个, [2]放第三个, ……
    phase_mix_prepare_list = [list(item) for item in zip(*phaseRads_list)]
    # PS-complex
    phase_mix_list = []
    phaseBit_mix_list = []
    phaseBitDeg_mix_list = []
    for phase_mix_prepare in phase_mix_prepare_list:
        phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_complex_beam_n(phase_mix_prepare, bit_num)
        phase_mix_list.append(phase_mix)
        phaseBit_mix_list.append(phaseBit_mix)
        phaseBitDeg_mix_list.append(phaseBitDeg_mix)
    return phase_mix_list, phaseBit_mix_list, phaseBitDeg_mix_list


# 多波束合成: QE
def multi_beam_qe(phaseBits_list):
    # 整理数据, 将所有[0]放第一个, [1]放第二个, [2]放第三个, ……
    phaseBit_mix_prepare_list = [list(item) for item in zip(*phaseBits_list)]
    # QE
    phaseBit_mix_list = []
    for phaseBit_mix_prepare in phaseBit_mix_prepare_list:
        phaseBit_mix, random_indices = qe_beam_N(phaseBit_mix_prepare)
        phaseBit_mix_list.append(phaseBit_mix)
    return phaseBit_mix_list


# 计算fitness组
def get_fitness(bit_num, beam_num, phaseBit, phaseBit_history, phaseBit_target):
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)
    fitness_function.phaseBitTarget = 1
    fitness_function.phaseBitHistoryList = phaseBit_history
    fitness_function.fitness_l4_map = {0: 2048, 1: 2048}
    fitness_function.init_main_lope_pos(phaseBit_target)
    # 前置计算, 计算pattern和peaks
    pattern, x, y = phaseBit_2_pattern(phaseBit, bit_num)
    peaks = pattern_2_peaks(pattern)
    #
    # 计算 L0
    l0 = fitness_function.fitness_l0(peaks, pattern)
    # 计算 L1
    l1 = fitness_function.fitness_l1(peaks)
    # 计算 L2
    l2 = fitness_function.fitness_l2(phaseBit)
    # 计算 L3
    l3 = fitness_function.fitness_l3(phaseBit)
    # 计算 L4
    l4 = fitness_function.fitness_l4(phaseBit)
    # 计算 L5
    l5 = fitness_function.fitness_l5(phaseBit, 1)
    # 计算 L
    fitness = fitness_function.fitness(phaseBit)
    logger.info("l=%f, l0=%d, l1=%f, l2=%s, l3=%s, l4=%s, l5=%d"
                % (fitness, l0, l1, l2, l3, l4, l5))
    return fitness, l0, l1, l2, l3, l4, l5


# 扫描的主程序
def main_beam_trace_bpso(path_pre, traces, phaseBit_list, bit_num, beam_num, bit_map, phaseBit_list_ps, phaseBit_list_qe):
    # 参数设置
    phaseBit_bpso_list = []                       # 优化后码阵 phaseBit_list
    path_pre_origin = path_pre + "origin/"      # 图片保存路径: 优化前
    path_pre_bpso = path_pre + "bpso/"              # 图片保存路径: 优化后
    path_pre_compare = path_pre + "compare/"    # 图片保存路径: 比较
    path_pre_ps = path_pre + "ps/"              # 图片保存路径: PS
    path_pre_qe = path_pre + "qe/"    # 图片保存路径: QE
    #
    fitness_function = FitnessMultiBeamTrace(bit_num, beam_num)
    bpso = BPSO(objective_function=fitness_function.fitness, num_particles=30, dimensions=(64, 64))
    rows, cols = phaseBit_list[0].shape
    # main-loop
    for i in range(len(traces)):
        phaseBit = phaseBit_list[i]
        phaseBit_ps = phaseBit_list_ps[i]
        phaseBit_qe = phaseBit_list_qe[i]
        if i > 0:
            phaseBit_ps_before = phaseBit_list_ps[i-1]
            phaseBit_qe_before = phaseBit_list_qe[i-1]
        else:
            phaseBit_ps_before = phaseBit_ps
            phaseBit_qe_before = phaseBit_qe
        logger.info("[BPSO Optimization main loop] i=%d(%d)" % (i, len(phaseBit_list)))
        # BPSO 寻优
        # 初始化Fitness方法, 主瓣位置, 历史码阵等
        fitness_function.init_main_lope_pos(phaseBit)
        fitness_function.phaseBitHistoryList = phaseBit_bpso_list
        fitness_function.fitness_l4_map = bit_map
        # BPSO 寻优
        phaseBit_bpso, best_fitness, phaseBit_bpso_fitness_history, best_individual_history \
            = bpso.optimize(30, phaseBit_ps_before, phaseBit_ps, phaseBit_qe_before, phaseBit_qe)
        phaseBit_bpso_list.append(phaseBit_bpso)
        # 计算方向图
        pattern_xyz, x, y = phaseBit_2_pattern(phaseBit, bit_num)
        pattern_xyz_bpso, x_bpso, y_bpso = phaseBit_2_pattern(phaseBit_bpso, bit_num)
        pattern_xyz_ps, x_ps, y_ps = phaseBit_2_pattern(phaseBit_ps, bit_num)
        pattern_xyz_qe, x_qe, y_qe = phaseBit_2_pattern(phaseBit_qe, bit_num)
        # # debug 画图
        # draw_img(phaseBit)
        # draw_img(phaseBit_bpso)
        # draw_img(np.abs(pattern_xyz))
        # draw_img(np.abs(pattern_xyz_bpso))
        #
        # 保存相位结果
        save_csv(phaseBit, path_pre_origin + "phaseBit_" + str(i) + ".csv")
        save_csv(phaseBit_bpso, path_pre_bpso + "phaseBit_bpso_" + str(i) + ".csv")
        save_csv(phaseBit_ps, path_pre_ps + "phaseBit_ps_" + str(i) + ".csv")
        save_csv(phaseBit_qe, path_pre_qe + "phaseBit_qe_" + str(i) + ".csv")
        # 设置前一时刻码阵, 开始时刻相等
        if i == 0:
            phaseBit_before = phaseBit
            phaseBit_bpso_before = phaseBit_bpso
        else:
            phaseBit_before = phaseBit_list[i - 1]
            phaseBit_bpso_before = phaseBit_bpso_list[i - 1]
        # 统计天线指标
        pg_list, psll, psll_coor, peaks = get_antenna_indicators_normalized_item(beam_num, pattern_xyz)
        pg_list_bpso, psll_bpso, psll_coor_bpso, peaks_bpso = get_antenna_indicators_normalized_item(beam_num, pattern_xyz_bpso)
        # 统计与前一个码阵的差别
        diff_phaseBit = compare_arrays(phaseBit_before, phaseBit)
        diff_phaseBit_bpso = compare_arrays(phaseBit_bpso_before, phaseBit_bpso)
        # 统计与前一个码阵的相差数量
        phaseBit_diff_count = np.sum(phaseBit != phaseBit_before)
        phaseBit_diff_count_percent = round((phaseBit_diff_count / (rows * cols)) * 100, 4)
        phaseBit_bpso_diff_count = np.sum(phaseBit_bpso != phaseBit_bpso_before)
        phaseBit_bpso_diff_count_percent = round((phaseBit_bpso_diff_count / (rows * cols)) * 100, 4)
        # 计算 fitness
        fitness, l0, l1, l2, l3, l4, l5 = get_fitness(bit_num, beam_num, phaseBit, phaseBit_list, phaseBit)
        fitness_bpso, l0_bpso, l1_bpso, l2_bpso, l3_bpso, l4_bpso, l5_bpso = get_fitness(bit_num, beam_num, phaseBit_bpso, phaseBit_bpso_list, phaseBit)
        fitness_ps, l0_ps, l1_ps, l2_ps, l3_ps, l4_ps, l5_ps = get_fitness(bit_num, beam_num, phaseBit_ps, phaseBit_list_ps, phaseBit)
        fitness_qe, l0_qe, l1_qe, l2_qe, l3_qe, l4_qe, l5_qe = get_fitness(bit_num, beam_num, phaseBit_qe, phaseBit_list_qe, phaseBit)
        # 保存图片
        img_text = "(theta,phi)=" + str(traces[i][0]) \
                   + ", count-diff=" + str(phaseBit_diff_count) \
                   + "(" + str(phaseBit_diff_count_percent) + ")" \
                   + ", count-diff(BPSO)=" + str(phaseBit_bpso_diff_count) \
                   + "(" + str(phaseBit_bpso_diff_count_percent) + ")" \
                   + ", PSLL=" + str(psll) + ", PSLL(BPSO)=" + str(psll_bpso)
                   # + ", PG=" + str(pg_list) + ", PG(BPSO)=" + str(pg_list_bpso) \
        save_images_3x2(path_pre_compare + "phaseBit_pattern" + str(i) + ".jpg",
                        phaseBit, f"phaseBit", np.abs(pattern_xyz), f"pattern", diff_phaseBit, f"diff(phaseBit)",
                        phaseBit_bpso, f"phaseBit(BPSO)", np.abs(pattern_xyz_bpso), f"pattern(BPSO)", diff_phaseBit_bpso,
                        f"diff(phaseBit(BPSO))",
                        img_text)
        # 保存遗传算法优化结果
        best_fitness_history = phaseBit_bpso_fitness_history
        save_line_chart(path_pre_bpso + "best_fitness_history" + str(i) + ".jpg", best_fitness_history,
                        "best_fitness_(theta,phi)=" + str(traces[i]), "iteration", "best_fitness", "fitness")
        save_csv([[item] for item in best_fitness_history], path_pre_bpso + "best_fitness_history" + str(i) + ".csv")
        # 保存指标
        result_csv = {
            "(theta,phi)": traces[i],
            "pg_list": pg_list,
            "pg_list_BPSO": pg_list_bpso,
            "psll": psll,
            "psll_BPSO": psll_bpso,
            "psll_coor": psll_coor,
            "psll_coor_BPSO": psll_coor_bpso,
            "phase_count_diff": phaseBit_diff_count,
            "phase_count_diff_BPSO": phaseBit_bpso_diff_count,
            "phase_count_percent_diff": phaseBit_diff_count_percent,
            "phase_count_percent_diff_BPSO": phaseBit_bpso_diff_count_percent,
            "fitness": fitness,
            "fitness_bpso": fitness_bpso,
            "fitness_ps": fitness_ps,
            "fitness_qe": fitness_qe,
            "l0": l0,
            "l0_bpso": l0_bpso,
            "l0_ps": l0_ps,
            "l0_qe": l0_qe,
            "l1": l1,
            "l1_bpso": l1_bpso,
            "l1_ps": l1_ps,
            "l1_qe": l1_qe,
            "l2": l2,
            "l2_bpso": l2_bpso,
            "l2_ps": l2_ps,
            "l2_qe": l2_qe,
            "l3": l3,
            "l3_bpso": l3_bpso,
            "l3_ps": l3_ps,
            "l3_qe": l3_qe,
            "l4": l4,
            "l4_bpso": l4_bpso,
            "l4_ps": l4_ps,
            "l4_qe": l4_qe,
            "l5": l5,
            "l5_bpso": l5_bpso,
            "l5_ps": l5_ps,
            "l5_qe": l5_qe,
        }
        save_to_csv([result_csv], path_pre_compare + "result.csv")



# 主程序: 多波束 -- 波束太多时使用
def main_beam_multi_trace_bpso_beams(path_pre, bit_num, bit_map):
    logger.info("main_beam_multi_trace_bpso_beams: bit_num=%d" % (bit_num))
    traces_list = []
    phaseRads_list = []
    phaseBit_list = []
    # 生成轨迹指向 & 码阵
    # phis = np.arange(0, 360, 90)          # 4波束
    # phis = np.arange(0, 360, 45)  # 8波束
    phis = np.arange(0, 360, 22.5)          # 16波束
    # phis = np.arange(0, 360, 11.25)         # 32波束
    logger.info("phis:%s" % phis)
    for phi in phis:
        print(phi)
        # 生成轨迹指向
        traces = create_scan_traces_theta(0, 60, phi)   # theta方向, (theta, phi): (30,  phi) --> (35,  phi)
        # 生成码阵
        phaseRads = traces_2_phaseRads(traces)
        # 码阵bit化
        phaseBit = phases_rad_2_bit(phaseRads, bit_num)
        logger.info("traces:%s" % (traces))
        traces_list.append(traces)
        phaseRads_list.append(phaseRads)
        phaseBit_list.append(phaseBit)
    # 合成多波束: PS-complex
    phase_mix_list_ps, phaseBit_mix_list_ps, phaseBitDeg_mix_list_ps = multi_beam_ps_complex(phaseRads_list, bit_num)
    # 合成多波束: QE
    phaseBit_mix_list_qe = multi_beam_qe(phaseBit_list)
    # 相邻码阵 BPSO 优化
    traces_mix_list = [list(item) for item in zip(*traces_list)]
    main_beam_trace_bpso(path_pre, traces_mix_list, phaseBit_mix_list_qe, bit_num, len(phaseBit_list), bit_map,
                    phaseBit_mix_list_ps, phaseBit_mix_list_qe)





if __name__ == '__main__':
    logger.info("multi-beam-trace-BPSO")
    # 主函数
    bit_map = {0: 2048, 1: 2048}
    main_beam_multi_trace_bpso_beams("../files/multi-beam-trace/beam2/multi-group-BPSO/2025-04-04-0903-16beam/",
                                   1, bit_map)
